import os
import sys
import argparse
import traceback
from pathlib import Path
from tqdm import tqdm  # 进度条库
import fitz  # PyMuPDF
import chardet  # 字符编码检测

def detect_encoding(text):
    """检测文本的编码"""
    result = chardet.detect(text)
    return result['encoding']

def convert_pdf_to_txt_pymupdf(pdf_path):
    """
    使用PyMuPDF (fitz) 将PDF文件转换为TXT
    PyMuPDF通常比PyPDF2有更好的文本提取能力
    
    参数:
    pdf_path (str): PDF文件的路径
    
    返回:
    str: 创建的TXT文件的路径
    """
    # 创建与原文件同名但扩展名为.txt的输出文件路径
    output_path = os.path.splitext(pdf_path)[0] + '.txt'
    
    try:
        # 打开PDF文件
        doc = fitz.open(pdf_path)
        
        # 检查PDF是否被加密
        if doc.is_encrypted:
            try:
                # 尝试空密码
                success = doc.authenticate("")
                if not success:
                    print(f"  错误: PDF文件被加密且需要密码 - {pdf_path}")
                    doc.close()
                    return None
            except Exception as e:
                print(f"  错误: 处理加密PDF时出错 - {str(e)}")
                doc.close()
                return None
        
        # 获取页数
        num_pages = len(doc)
        print(f"  PDF页数: {num_pages}")
        
        if num_pages == 0:
            print(f"  警告: PDF没有页面 - {pdf_path}")
            doc.close()
            return None
        
        # 提取文本内容
        text = ""
        for page_num in range(num_pages):
            page = doc[page_num]
            page_text = page.get_text()
            if page_text:
                text += page_text + "\n\n"
            else:
                print(f"  警告: 页面 {page_num+1} 未提取到文本")
        
        # 关闭文档
        doc.close()
        
        if not text.strip():
            print(f"  警告: 未能从PDF中提取到任何文本 - {pdf_path}")
            print(f"  这可能是一个扫描PDF，需要OCR处理")
            return None
        
        # 写入TXT文件
        try:
            with open(output_path, 'w', encoding='utf-8') as txt_file:
                txt_file.write(text)
        except UnicodeEncodeError:
            # 如果UTF-8编码失败，尝试检测和使用其他编码
            encoding = detect_encoding(text.encode('utf-8', errors='replace'))
            print(f"  文本编码检测为: {encoding}")
            with open(output_path, 'w', encoding=encoding or 'utf-8', errors='replace') as txt_file:
                txt_file.write(text)
        
        print(f"  成功: {pdf_path} -> {output_path}")
        return output_path
        
    except Exception as e:
        print(f"  错误转换 {pdf_path}: {str(e)}")
        print(f"  详细错误: {traceback.format_exc()}")
        return None

def batch_convert_pdfs(input_dir):
    """批量转换目录中的所有PDF文件为TXT文件"""
    # 确保使用绝对路径
    input_path = Path(os.path.abspath(input_dir))
    
    print(f"使用路径: {input_path}")
    
    # 检查目录是否存在
    if not input_path.exists():
        print(f"错误: 目录不存在 - {input_path}")
        return 0, 0, 0
        
    if not input_path.is_dir():
        print(f"错误: 不是一个目录 - {input_path}")
        return 0, 0, 0
    
    # 查找所有PDF文件（不区分大小写）
    pdf_files = []
    for file in input_path.iterdir():
        if file.is_file() and file.suffix.lower() == '.pdf':
            pdf_files.append(file)
    
    total_files = len(pdf_files)
    
    if total_files == 0:
        print(f"警告: 在目录 {input_path} 中没有找到PDF文件")
        print(f"目录中的文件:")
        for file in input_path.iterdir():
            if file.is_file():
                print(f"  - {file.name} ({file.suffix})")
        return 0, 0, 0
    
    print(f"找到 {total_files} 个PDF文件:")
    for pdf_file in pdf_files:
        print(f"  - {pdf_file.name}")
    
    # 初始化计数器
    success_count = 0
    failure_count = 0
    
    # 使用进度条处理文件
    for pdf_file in tqdm(pdf_files, desc="处理PDF文件", unit="文件"):
        print(f"\n处理文件: {pdf_file.name}")
        
        # 转换PDF到TXT
        result = convert_pdf_to_txt_pymupdf(str(pdf_file))
        
        if result:
            success_count += 1
        else:
            failure_count += 1
    
    return success_count, failure_count, total_files

def main():
    # 创建命令行参数解析器
    parser = argparse.ArgumentParser(description='高级批量PDF到TXT转换工具（支持更多PDF格式）')
    parser.add_argument('input_dir', help='包含PDF文件的目录')
    
    # 解析命令行参数
    args = parser.parse_args()
    
    # 打印系统信息
    print(f"当前工作目录: {os.getcwd()}")
    print(f"Python 版本: {sys.version}")
    print(f"PyMuPDF 版本: {fitz.version[0]}")
    print(f"处理目录: {args.input_dir}\n")
    
    # 执行批量转换
    success, failure, total = batch_convert_pdfs(args.input_dir)
    
    # 打印转换结果摘要
    print(f"\n转换完成：")
    print(f"找到的PDF文件: {total}")
    print(f"成功转换: {success} 文件")
    print(f"失败转换: {failure} 文件")
    
    if failure > 0:
        print("\n对于失败的文件，可能的原因包括:")
        print("1. PDF文件受密码保护")
        print("2. PDF文件已损坏")
        print("3. PDF文件是扫描件，仅包含图像而没有文本")
        print("4. PDF使用了特殊的字体或编码")
        print("\n如果需要处理扫描PDF，可能需要使用OCR工具，如：")
        print("- pytesseract + pdf2image")
        print("- ocrmypdf")

if __name__ == "__main__":
    main()